package org.apache.lucene.analysis.pinyin.lucene6;

import java.io.IOException;
import org.apache.lucene.analysis.TokenFilter;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.pinyin.utils.StringUtils;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.OffsetAttribute;
import org.apache.lucene.analysis.tokenattributes.PositionIncrementAttribute;
import org.apache.lucene.analysis.tokenattributes.TypeAttribute;
//import org.apache.lucene.analysis.util.CharacterUtils;

/**
 * 对转换后的拼音进行EdgeNGram处理的TokenFilter
 * @author Lanxiaowei
 */
public class PinyinEdgeNGramTokenFilter extends TokenFilter {
    private final int minGram;
    private final int maxGram;
    private final boolean nGramChinese;
    private final boolean nGramNumber;
    private CharTermAttribute termAtt;
    private OffsetAttribute offsetAtt;
    private PositionIncrementAttribute posIncrAtt;
    private TypeAttribute typeAtt;
    private char[] curTermBuffer;
    private int curTermLength;
    private int curGramSize;
    private int tokStart;
    private PinyinEdgeNGramTokenFilter.Side side;
    public static final PinyinEdgeNGramTokenFilter.Side DEFAULT_SIDE;

    static {
        DEFAULT_SIDE = PinyinEdgeNGramTokenFilter.Side.FRONT;
    }

    public PinyinEdgeNGramTokenFilter(TokenStream input) {
        this(input, "front", 2, 10, false);
    }

    public PinyinEdgeNGramTokenFilter(TokenStream input, String side) {
        this(input, side, 2, 10, false);
    }

    public PinyinEdgeNGramTokenFilter(TokenStream input, String side, int maxGram) {
        this(input, side, 2, maxGram, false);
    }

    public PinyinEdgeNGramTokenFilter(TokenStream input, String side, int minGram, int maxGram) {
        this(input, side, minGram, maxGram, false, false);
    }

    public PinyinEdgeNGramTokenFilter(TokenStream input, String side, int minGram, int maxGram, boolean nGramChinese) {
        this(input, side, minGram, maxGram, nGramChinese, false);
    }

    public PinyinEdgeNGramTokenFilter(TokenStream input, String side, int minGram, int maxGram, boolean nGramChinese,
        boolean nGramNumber) {
        super(input);
        if (minGram < 1) {
            throw new IllegalArgumentException("minGram must be greater than zero");
        } else if (minGram > maxGram) {
            throw new IllegalArgumentException("minGram must not be greater than maxGram");
        } else if (side != null && (side.equalsIgnoreCase("front") || side.equalsIgnoreCase("back"))) {
            this.termAtt = (CharTermAttribute) this.addAttribute(CharTermAttribute.class);
            this.offsetAtt = (OffsetAttribute) this.addAttribute(OffsetAttribute.class);
            this.posIncrAtt = (PositionIncrementAttribute) this.addAttribute(PositionIncrementAttribute.class);
            this.typeAtt = (TypeAttribute) this.addAttribute(TypeAttribute.class);
            this.minGram = minGram;
            this.maxGram = maxGram;
            this.nGramChinese = nGramChinese;
            this.nGramNumber = nGramNumber;
            this.side = PinyinEdgeNGramTokenFilter.Side.getSide(side);
        } else {
            throw new IllegalArgumentException("side must be either front or back");
        }
    }

    public final boolean incrementToken() throws IOException {
        while (true) {
            if (this.curTermBuffer == null) {
                if (!this.input.incrementToken()) {
                    return false;
                }

                String start = this.typeAtt.type();
                if (start != null && "normal_word".equals(start)) {
                    return true;
                }

                if (start != null && "numeric_original".equals(start)) {
                    return true;
                }

                if (start != null && "chinese_original".equals(start)) {
                    return true;
                }

                if (!this.nGramNumber && StringUtils.isNumeric(this.termAtt.toString())) {
                    return true;
                }

                if (!this.nGramChinese && StringUtils.containsChinese(this.termAtt.toString())) {
                    return true;
                }

                this.curTermBuffer = this.termAtt.buffer().clone();
                this.curTermLength = this.termAtt.length();
                this.curGramSize = this.minGram;
                this.tokStart = this.offsetAtt.startOffset();
            }

            if (this.curTermLength >= this.minGram && this.curGramSize <= this.maxGram) {
                if (this.curGramSize >= this.curTermLength) {
                    this.clearAttributes();
                    this.offsetAtt.setOffset(this.tokStart + 0, this.tokStart + this.curTermLength);
                    this.termAtt.copyBuffer(this.curTermBuffer, 0, this.curTermLength);
                    this.posIncrAtt.setPositionIncrement(0);
                    this.curTermBuffer = null;
                    return true;
                }

                int start1 = this.side == PinyinEdgeNGramTokenFilter.Side.FRONT ? 0 : this.curTermLength - this.curGramSize;
                int end = start1 + this.curGramSize;
                this.clearAttributes();
                this.offsetAtt.setOffset(this.tokStart + start1, this.tokStart + end);
                this.termAtt.copyBuffer(this.curTermBuffer, start1, this.curGramSize);
                this.posIncrAtt.setPositionIncrement(0);
                ++this.curGramSize;
                return true;
            }

            this.curTermBuffer = null;
        }
    }

    public void reset() throws IOException {
        super.reset();
        this.curTermBuffer = null;
    }

    public enum Side {
        FRONT {
            public String getLabel() {
                return "front";
            }
        },
        BACK {
            public String getLabel() {
                return "back";
            }
        };

        Side() {
        }

        public abstract String getLabel();

        public static PinyinEdgeNGramTokenFilter.Side getSide(String sideName) {
            return FRONT.getLabel().equalsIgnoreCase(sideName) ? FRONT : (BACK.getLabel().equalsIgnoreCase(sideName) ? BACK : null);
        }
    }
}